package com.springboot.demo.lock.redis;

import java.util.Collections;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import javax.annotation.PostConstruct;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Service;

/**
 * redis限流
 *
 * @author supu
 * @date 2020-03-25 15:44
 **/
@Service
public class RedisLimit {
    private Map<Integer, Integer> resetNumberMap = new ConcurrentHashMap<>(16);

    private final DefaultRedisScript<Long> redisLimitScript = new DefaultRedisScript<>();
    private final DefaultRedisScript<Integer> redisGetNumScript = new DefaultRedisScript<>();

    private static final Integer DOUBLE_DIGIT = 2;

    @Autowired
    private RedisTemplate redisTemplate;

    @PostConstruct
    private void initData() {
        // 初始化位数数据
        resetNumberMap.put(1, 9);
        resetNumberMap.put(2, 99);
        resetNumberMap.put(3, 999);
        resetNumberMap.put(4, 9999);
        resetNumberMap.put(5, 99999);
        resetNumberMap.put(6, 999999);
        resetNumberMap.put(7, 9999999);
        resetNumberMap.put(8, 99999999);

        // 初始化DefaultRedisScript对象
        redisLimitScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("/lua/redis_limit.lua")));
        redisLimitScript.setResultType(Long.class);

        // 初始化DefaultRedisScript对象
        redisGetNumScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("/lua/redis_get_number.lua")));
        redisGetNumScript.setResultType(Integer.class);
    }

    public Long redisLimit(String key, int limit) {
        return (Long) redisTemplate.execute(redisLimitScript, Collections.singletonList(key), limit);
    }

    public Integer redisGetNum(String key, int resetNum) {
        return (Integer) redisTemplate.execute(redisGetNumScript, Collections.singletonList(key), resetNum);
    }

    /**
     * 根据key获取指定位数的字符串
     *
     * @param key 缓存key
     * @param num 字符串位数
     * @return
     */
    public String getDigitsString(String key, int num) {
        int number = redisGetNum(key, resetNumberMap.get(num));
        StringBuffer stringBuffer = new StringBuffer();
        if (num > 1 && number < DOUBLE_DIGIT) {
            // 个位数时前面补0
            stringBuffer.append("0");
        }
        stringBuffer.append(number);
        return stringBuffer.toString();
    }
}
